
import subprocess
import sys
import difflib
import os

def get_command_output(command, exclude_lines=0, capture_stderr=False):
    """
    执行命令并获取输出，可以选择捕获 stderr 并排除前几行
    """
    try:
        # 设置输出目标
        stdout_dest = subprocess.PIPE
        stderr_dest = subprocess.STDOUT if capture_stderr else subprocess.PIPE
        
        # 执行命令
        result = subprocess.run(
            command, 
            shell=True, 
            stdout=stdout_dest, 
            stderr=stderr_dest,
            text=True
        )
        
        # 处理不同情况下的输出
        if capture_stderr:
            # 当捕获stderr时，输出在stdout中（因为stderr被重定向）
            output = result.stdout
        else:
            # 正常情况下使用stdout
            output = result.stdout
            
        # 如果没有任何输出，返回空列表
        if output is None or output.strip() == "":
            return []
        
        # 分割输出行
        lines = output.splitlines()
        
        # 排除前几行
        if exclude_lines > 0 and len(lines) > exclude_lines:
            return lines[exclude_lines:]
        else:
            return lines
            
    except subprocess.CalledProcessError as e:
        print(f"命令执行失败: {command}")
        print(f"返回码: {e.returncode}")
        print(f"错误信息: {e.stderr or e.stdout}")
        return []
    except Exception as e:
        print(f"执行命令时出错: {command}")
        print(f"错误信息: {str(e)}")
        return []

def main():
    if len(sys.argv) < 2:
        print("请提供测试文件名 (例如: test0.c)")
        sys.exit(1)
    
    test_file = sys.argv[1]
    base_name = os.path.splitext(os.path.basename(test_file))[0].replace("test", "")
    ans_file = f"ans{base_name}"
    
    # 获取 cargo run 的输出（排除前两行，捕获 stderr）
    cargo_output = get_command_output(f"cargo run {test_file}", exclude_lines=2, capture_stderr=True)
    
    # 获取 ans 文件的输出
    if not os.path.exists(ans_file):
        print(f"错误: ans 文件 {ans_file} 不存在")
        sys.exit(1)
    
    ans_output = get_command_output(f"cat {ans_file}")
    
    # 添加详细的调试信息
    print("\n" + "="*50)
    print(f"测试文件: {test_file}")
    print(f"答案文件: {ans_file}")
    print(f" cargo run 输出行数: {len(cargo_output)}")
    print(f" ans 文件输出行数: {len(ans_output)}")
    print("="*50 + "\n")
    
    # 如果两个输出都为空，特殊处理
    if not cargo_output and not ans_output:
        print("\033[92m两者输出均为空，被视为一致!\033[0m")
        sys.exit(0)
    
    # 比较两个输出
    diff = difflib.unified_diff(
        cargo_output, 
        ans_output,
        fromfile=f"cargo run {test_file}",
        tofile=ans_file,
        lineterm=''
    )
    
    diff_list = list(diff)
    
    if len(diff_list) == 0:
        print("\033[92m输出完全一致!\033[0m")  # 绿色
    else:
        print("\033[91m发现差异:\033[0m")  # 红色
        for line in diff_list:
            if line.startswith('+'):
                print(f"\033[92m{line}\033[0m")  # 绿色显示新增行
            elif line.startswith('-'):
                print(f"\033[91m{line}\033[0m")  # 红色显示删除行
            else:
                print(line)

if __name__ == "__main__":
    main()
